// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/exchange_queue.h>
#include <algorithm>

namespace kumo::pollux::exec {
    SerializedPage::SerializedPage(
        std::unique_ptr<melon::IOBuf> iobuf,
        std::function<void(melon::IOBuf &)> onDestructionCb,
        std::optional<int64_t> numRows)
        : iobuf_(std::move(iobuf)),
          iobufBytes_(chainBytes(*iobuf_.get())),
          numRows_(numRows),
          onDestructionCb_(onDestructionCb) {
        POLLUX_CHECK_NOT_NULL(iobuf_);
        for (auto &buf: *iobuf_) {
            int32_t bufSize = buf.size();
            ranges_.push_back(ByteRange{
                const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(buf.data())),
                bufSize,
                0
            });
        }
    }

    SerializedPage::~SerializedPage() {
        if (onDestructionCb_) {
            onDestructionCb_(*iobuf_.get());
        }
    }

    std::unique_ptr<ByteInputStream> SerializedPage::prepareStreamForDeserialize() {
        return std::make_unique<BufferInputStream>(std::move(ranges_));
    }

    void ExchangeQueue::noMoreSources() {
        std::vector<ContinuePromise> promises; {
            std::lock_guard<std::mutex> l(mutex_);
            noMoreSources_ = true;
            promises = checkCompleteLocked();
        }
        clearPromises(promises);
    }

    void ExchangeQueue::close() {
        std::vector<ContinuePromise> promises; {
            std::lock_guard<std::mutex> l(mutex_);
            promises = closeLocked();
        }
        clearPromises(promises);
    }

    int64_t ExchangeQueue::minOutputBatchBytesLocked() const {
        // always allow to unblock when at end
        if (atEnd_) {
            return 0;
        }
        // At most 1% of received bytes so far to minimize latency for small exchanges
        return std::min<int64_t>(minOutputBatchBytes_, receivedBytes_ / 100);
    }

    void ExchangeQueue::enqueueLocked(
        std::unique_ptr<SerializedPage> &&page,
        std::vector<ContinuePromise> &promises) {
        if (page == nullptr) {
            ++numCompleted_;
            auto completedPromises = checkCompleteLocked();
            promises.reserve(promises.size() + completedPromises.size());
            for (auto &promise: completedPromises) {
                promises.push_back(std::move(promise));
            }
            return;
        }

        totalBytes_ += page->size();
        if (peakBytes_ < totalBytes_) {
            peakBytes_ = totalBytes_;
        }

        ++receivedPages_;
        receivedBytes_ += page->size();

        queue_.push_back(std::move(page));
        const auto minBatchSize = minOutputBatchBytesLocked();
        while (!promises_.empty()) {
            POLLUX_CHECK_LE(promises_.size(), numberOfConsumers_);
            const int32_t unblockedConsumers = numberOfConsumers_ - promises_.size();
            const int64_t unasignedBytes =
                    totalBytes_ - unblockedConsumers * minBatchSize;
            if (unasignedBytes < minBatchSize) {
                break;
            }
            // Resume one of the waiting drivers.
            auto it = promises_.begin();
            promises.push_back(std::move(it->second));
            promises_.erase(it);
        }
    }

    void ExchangeQueue::addPromiseLocked(
        int consumerId,
        ContinueFuture *future,
        ContinuePromise *stalePromise) {
        ContinuePromise promise{"ExchangeQueue::dequeue"};
        *future = promise.getSemiFuture();
        auto it = promises_.find(consumerId);
        if (it != promises_.end()) {
            // resolve stale promises outside the lock to avoid broken promises
            *stalePromise = std::move(it->second);
            it->second = std::move(promise);
        } else {
            promises_[consumerId] = std::move(promise);
        }
        POLLUX_CHECK_LE(promises_.size(), numberOfConsumers_);
    }

    std::vector<std::unique_ptr<SerializedPage> > ExchangeQueue::dequeueLocked(
        int consumerId,
        uint32_t maxBytes,
        bool *atEnd,
        ContinueFuture *future,
        ContinuePromise *stalePromise) {
        POLLUX_CHECK_NOT_NULL(future);
        if (!error_.empty()) {
            *atEnd = true;
            POLLUX_FAIL(error_);
        }

        *atEnd = false;

        // If we don't have enough bytes to return, we wait for more data to be
        // available
        if (totalBytes_ < minOutputBatchBytesLocked()) {
            addPromiseLocked(consumerId, future, stalePromise);
            return {};
        }

        std::vector<std::unique_ptr<SerializedPage> > pages;
        uint32_t pageBytes = 0;
        for (;;) {
            if (queue_.empty()) {
                if (atEnd_) {
                    *atEnd = true;
                } else if (pages.empty()) {
                    addPromiseLocked(consumerId, future, stalePromise);
                }
                return pages;
            }

            if (pageBytes > 0 && pageBytes + queue_.front()->size() > maxBytes) {
                return pages;
            }

            pages.emplace_back(std::move(queue_.front()));
            queue_.pop_front();
            pageBytes += pages.back()->size();
            totalBytes_ -= pages.back()->size();
        }

        POLLUX_UNREACHABLE();
    }

    void ExchangeQueue::setError(const std::string &error) {
        std::vector<ContinuePromise> promises; {
            std::lock_guard<std::mutex> l(mutex_);
            if (!error_.empty()) {
                return;
            }
            error_ = error;
            atEnd_ = true;
            // NOTE: clear the serialized page queue as we won't consume from an
            // errored queue.
            queue_.clear();
            promises = clearAllPromisesLocked();
        }
        clearPromises(promises);
    }
} // namespace kumo::pollux::exec
